import os
import time

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.optim.lr_scheduler import MultiStepLR
from torchnet.meter import ConfusionMeter
import matplotlib.pyplot as plt

from networks import create_resnet50, create_my_cnn
from pickle_utils import save_variable


# 以下函数用于计算模型的分类准确率以及混淆矩阵
def evaluate(net, data_loader, num_classes, device):
    correct = 0
    total = 0
    cm = ConfusionMeter(num_classes)
    with torch.no_grad():
        for data in data_loader:
            images, labels = data[0].to(device), data[1].to(device)
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            cm.add(predicted, labels)

        acc = correct / total
    return acc, cm


def train(net, optimizer, loss_fn, num_epoch, data_loader, device, num_classes=20, test_loader=None, scheduler=None):
    # 训练过程中会在每一轮次后计算训练集和测试集准确率，和loss一并输出
    # 如果不指定测试集test_loader，则每一轮次后只计算训练集准确率并输出
    # 最终返回epoch, loss, train_acc, test_acc四个列表，代表其每一轮次变化，用于曲线绘制
    net.train()
    epochs = []
    losses = []
    train_accs = []
    test_accs = []
    best_loss = 9999
    best_epoch = 0
    train_begin = time.time()
    for epoch in range(num_epoch):
        running_loss = 0.0
        begin = time.time()
        for i, data in enumerate(data_loader):
            inputs, labels = data[0].to(device), data[1].to(device)

            optimizer.zero_grad()

            outputs = net(inputs)

            loss = loss_fn(outputs, labels)
            loss.backward()

            optimizer.step()

            running_loss += loss.item()

        if scheduler is not None:
            scheduler.step()

        end = time.time()
        epoch_time = end - begin
        step_time = epoch_time / len(data_loader)

        epoch_loss = running_loss / len(data_loader)

        epochs.append(epoch)
        losses.append(epoch_loss)

        net.eval()
        print('epoch %d :' % (epoch + 1))
        print('loss = %.3f' % epoch_loss)
        train_acc, _ = evaluate(net=net,
                                data_loader=data_loader,
                                num_classes=num_classes,
                                device=device)
        train_accs.append(train_acc)
        if test_loader is not None:
            test_acc, _ = evaluate(net=net,
                                   data_loader=test_loader,
                                   num_classes=num_classes,
                                   device=device)
            # net.train()
            print('train_acc = %.2f %%  test_acc = %.2f %%'
                  % (train_acc * 100, test_acc * 100))
            test_accs.append(test_acc)
        else:
            print('train_acc = %.2f %%' % (train_acc * 100))
        print('epoch_time = %d min %d s  step_time = %.2f ms'
              % (epoch_time / 60, epoch_time % 60, step_time * 1000))
        print('-' * 50)
        net.train()

        if epoch_loss < best_loss:
            best_loss = epoch_loss
            best_epoch = epoch
            best_model_wts = net.state_dict()

        if (epoch + 1) % 10 == 0:
            # 保存状态词典
            os.makedirs('checkpoints', exist_ok=True)
            state_dict = net.state_dict()
            torch.save(state_dict, f'checkpoints/net-{epoch + 1}.pth')

        plt.cla()
        plt.xlabel('epoch')
        plt.ylabel('loss')
        plt.plot(epochs, losses)
        plt.savefig("loss.png")

    elapsed = time.time() - train_begin
    print('train finished in %d min %d s' % (elapsed / 60, elapsed % 60))
    print('best epoch:', best_epoch + 1)
    torch.save(best_model_wts, 'checkpoints/net-best.pth')

    return epochs, losses, train_accs, test_accs


if __name__ == '__main__':
    transform = transforms.Compose(  # 训练集数据变换
        [
            torchvision.transforms.Resize((224, 224), interpolation=2),
            torchvision.transforms.ColorJitter(brightness=0, contrast=0, saturation=0, hue=0),
            transforms.RandomResizedCrop(size=160, scale=(0.80, 1.0), ratio=(0.75, 1.3333), interpolation=3),
            torchvision.transforms.RandomHorizontalFlip(p=0.5),

            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406), [0.229, 0.224, 0.225])
        ]
    )
    transform_test = transforms.Compose(  # 测试集数据变换
        [
            torchvision.transforms.Resize((224, 224), interpolation=2),
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406), [0.229, 0.224, 0.225])
        ]
    )

    dataset_path = 'dataset/oregon_wildlife'
    train_dataset_path = os.path.join(dataset_path, 'train')
    test_dataset_path = os.path.join(dataset_path, 'test')
    train_set = torchvision.datasets.ImageFolder(root=train_dataset_path, transform=transform)
    test_set = torchvision.datasets.ImageFolder(root=test_dataset_path, transform=transform_test)

    train_loader = torch.utils.data.DataLoader(dataset=train_set,
                                               batch_size=36,
                                               shuffle=True,
                                               num_workers=0)
    test_loader = torch.utils.data.DataLoader(dataset=test_set,
                                              batch_size=36,
                                              shuffle=False,
                                              num_workers=0)

    device = torch.device("cuda:0")

    num_classes = len(train_set.classes)
    print('num classes:', num_classes)

    # ====================== 设置不同的网络 ==========================

    # # ===================== my cnn ====================
    # # 请设置输入尺寸为(160, 160)
    # net = create_my_cnn(num_classes=num_classes)  # 创建网络对象
    # net.to(device)
    # optimizer = optim.SGD(net.parameters(), lr=0.002, momentum=0.9, weight_decay=5e-4, nesterov=True)
    # num_epoch = 60  # 轮次60
    # scheduler = MultiStepLR(optimizer, milestones=[40], gamma=0.1)  # lr在40个epoch后衰减为0.0002

    # # ================== resnet, no pretrained ======================
    # # 请设置输入尺寸为(224, 224)
    # net = create_resnet50(num_classes=num_classes, pretrained=False)
    # net.to(device)
    # optimizer = optim.SGD(net.parameters(), lr=0.002, momentum=0.9, weight_decay=5e-4, nesterov=True)
    # num_epoch = 60  # 60个epoch
    # scheduler = MultiStepLR(optimizer, milestones=[40], gamma=0.1)  # lr在8个epoch后衰减为0.0002

    # ================== resnet, pretrained =================
    # 请设置输入尺寸为(224, 224)
    net = create_resnet50(num_classes=num_classes, pretrained=True)
    net.to(device)
    optimizer = optim.RMSprop(net.parameters(), lr=0.0001)  # RMSprop，学习率0.0001
    num_epoch = 12  # 12个epoch
    scheduler = MultiStepLR(optimizer, milestones=[8], gamma=0.1)  # lr在8个epoch后衰减为0.00001

    xentropy = nn.CrossEntropyLoss()  # 交叉熵损失

    # 训练
    epochs, losses, train_accs, test_accs = train(net=net,
                                                  optimizer=optimizer,
                                                  loss_fn=xentropy,
                                                  num_epoch=num_epoch,
                                                  data_loader=train_loader,
                                                  device=device,
                                                  num_classes=num_classes,
                                                  test_loader=test_loader,
                                                  scheduler=scheduler
                                                  )

    net.eval()

    # 保存状态词典
    os.makedirs('checkpoints', exist_ok=True)
    state_dict = net.state_dict()
    torch.save(state_dict, 'checkpoints/net-last.pth')
    print('Done')

    # 保存损失、准确率变化
    save_variable(losses, 'losses.pkl')
    save_variable(train_accs, 'train_accs.pkl')
    save_variable(test_accs, 'test_accs.pkl')

    # 计算准确率和混淆矩阵
    train_acc, cm_train = evaluate(net=net,
                                   data_loader=train_loader,
                                   num_classes=num_classes,
                                   device=device)
    test_acc, cm = evaluate(net=net,
                            data_loader=test_loader,
                            num_classes=num_classes,
                            device=device)
    print('Training Accuracy: %.2f %%' % (100 * train_acc))
    print(cm_train.value())
    print('Test Accuracy: %.2f %%' % (100 * test_acc))
    print(cm.value())
    eval_info = {
        'train_acc': train_acc,
        'test_acc': test_acc,
        'cm_train': cm_train.value(),
        'cm_test': cm.value()
    }
    save_variable(eval_info, 'eval_info.pkl')

    # 绘制loss曲线
    plt.xlabel('epoch')
    plt.ylabel('loss')
    plt.plot(epochs, losses)
    plt.show()

    # 绘制并保存准确率曲线
    plt.xlabel('epoch')
    plt.ylabel('acc')
    plt.plot(epochs, train_accs, label="train")
    plt.plot(epochs, test_accs, label="test")
    plt.legend(loc='lower right')
    plt.savefig("acc.png")
